import torch
import time
import sys
from utils.util import *
from utils.save import *
from torch.nn import functional as F
def validate(args, val_loader, model, criterion, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top2 = AverageMeter()
    log = Log()
    model.eval()   
    end = time.time()

    # we may have ten d in data
    for i, (data, target, paths) in enumerate(val_loader):
        if args.gpu is not None:
            data = data.cuda()
            target = target.cuda()

        # compute output
        for idx, d in enumerate(data[0]):      # data [batchsize, 10_crop, 3, 448, 448]
            d = d.unsqueeze(0) # d [1, 3, 448, 448]
            output1, output2, output3, _ = model(d)
            output = output1 + output2 + 0.1 * output3

            # measure accuracy and record loss
            prec1, prec2 = accuracy(output, target, topk=(1, 2))

            top1.update(prec1[0], 1)
            top2.update(prec2[0], 1)
            print('Test <==> Img:{} No:{} Top1 {:.3f} Top2 {:.3f}'.format(i, idx, prec1.cpu().numpy()[0], prec2.cpu().numpy()[0]))

    print('Test Total <==> Top1 {:.3f}% Top2 {:.3f}%'.format(top1.avg, top2.avg))
    log.save_test_info(epoch, top1, top2)
    return top1.avg


# def validate_simple(args, val_loader, model, criterion, epoch):
#     batch_time = AverageMeter()
#     losses = AverageMeter()
#     top1 = AverageMeter()
#     top2 = AverageMeter()
#     log = Log()
#     model.eval()
#     end = time.time()
#
#     # we may have ten d in data
#     for i, (data, target, paths) in enumerate(val_loader):
#         if args.gpu is not None:
#             data = data.cuda()
#             target = target.cuda()
#
#         # compute output
#         for idx, d in enumerate(data):      # data [batchsize, 10_crop, 3, 448, 448]
#             d = d.unsqueeze(0) # d [1, 3, 448, 448]
#             output1, output2, output3, _ = model(d)
#             output = output1 + output2 + 0.1 * output3
#
#             # measure accuracy and record loss
#             prec1, prec2 = accuracy(output, target, topk=(1, 2))
#
#             top1.update(prec1[0], 1)
#             top2.update(prec2[0], 1)
#             print('DFL-CNN <==> Test <==> Img:{} Top1 {:.3f} Top2 {:.3f}'.format(i, prec1.cpu().numpy()[0], prec2.cpu().numpy()[0]))
#
#     print('DFL-CNN <==> Test Total <==> Top1 {:.3f}% Top2 {:.3f}%'.format(top1.avg, top2.avg))
#     log.save_test_info(epoch, top1, top2)
#     return top1.avg


def validate_simple(args, val_loader, model, criterion, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top2 = AverageMeter()
    log = Log()
    model.eval()
    end = time.time()
    target_list = []
    predict_list = []
    # we may have ten d in data
    for i, (data, target, paths) in enumerate(val_loader):
        if args.gpu is not None:
            data = data.cuda()
            target = target.cuda()

        # compute output
        for idx, d in enumerate(data):      # data [batchsize, 10_crop, 3, 448, 448]
            d = d.unsqueeze(0) # d [1, 3, 448, 448]
            output = model(d)
            # output = output1 + output2 + 0.1 * output3
            out_float = F.softmax(output).data.cpu().numpy()[0].argmax()
            target_int = target.item()
            target_list.append(out_float)
            predict_list.append(target_int)
            # measure accuracy and record loss
            prec = accuracy(output, target)

            top1.update(prec[0], 1)
            # top2.update(prec2[0], 1)
            print('Test <==> Img:{} Top1 {}'.format(i, prec[0].item()))
    from sklearn import metrics
    precision_score = metrics.precision_score(target_list, predict_list, average='macro')
    recall_score = metrics.recall_score(target_list, predict_list, average='macro')
    pre_score = metrics.precision_score(target_list, predict_list, labels=[0], average='macro')
    rec_score = metrics.recall_score(target_list, predict_list,labels=[0], average='macro')
    print('precision:',precision_score)
    print('recall:', recall_score)
    print('pre_score:',pre_score)
    print('rec_score:', rec_score)
    print('Test Total <==> Top1 {:.3f}% '.format(top1.avg[0]))
    with open('/home/ubuntu/data/log/log_effic.txt', 'a') as f:
        f.write('Integer Label,efficient-3,Feature 55')
        f.write('Test Total <==> epoch{} Top1 {:.3f}% {:.3f}% {:.3f}%\n'.format(epoch, top1.avg[0],pre_score,rec_score))

    return top1.avg

def muti_validate_simple(args, val_loader, model, criterion, epoch):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top2 = AverageMeter()
    log = Log()
    model.eval()
    end = time.time()

    a = val_loader.dataset.class_to_idx
    # we may have ten d in data
    for i, (data, target, paths) in enumerate(val_loader):
        class_list = [[int(list(a.keys())[list(a.values()).index(i)][j]) for i in target] for j in range(args.nclass)]
        class_list = torch.tensor(class_list)
        if args.gpu is not None:
            data = data.cuda()
            target = target.cuda()
            class_list = class_list.cuda()
        # compute output
        for idx, d in enumerate(data):      # data [batchsize, 10_crop, 3, 448, 448]
            d = d.unsqueeze(0) # d [1, 3, 448, 448]
            output = model(d)
            # output = output1 + output2 + 0.1 * output3

            # measure accuracy and record loss
            for k in range(len(class_list)):
                prec = accuracy(output[k], class_list[k])

                if k ==0:
                    top1.update(prec[0], 1)
                    print('Test <==> Img:{} Top1 {}'.format(i, prec[0].item()))
                if k == 1:
                    top2.update(prec[0], 1)
                    print('Test <==> Img:{} Top2 {}'.format(i, prec[0].item()))

    print('Test Total <==> Top1 {:.3f}% '.format(top1.avg[0]))
    print('Test Total <==> Top2 {:.3f}% '.format(top2.avg[0]))
    with open('/home/ubuntu/data/log/muti_effic.txt', 'a') as f:
        f.write('Integer Label,efficient,Feature 14')
        f.write('Test Total <==>Epoch{} Top1 {:.3f}%  Top2 {:.3f}%\n'.format(epoch, top1.avg[0], top2.avg[0]))

    return (top1.avg + top1.avg)/2.


